import os
import sys
import subprocess
import importlib

def check_and_install_dependencies():
    """检查并安装所需的依赖"""
    print("正在检查并安装必要的依赖...")
    
    # 安装 detectron2
    try:
        import detectron2
        print("Detectron2 已安装")
    except ImportError:
        print("正在安装 Detectron2...")
        subprocess.check_call([
            sys.executable, "-m", "pip", "install", 
            "git+https://github.com/facebookresearch/detectron2.git"
        ])
    
    # 安装 layoutparser 及其依赖
    try:
        import layoutparser
        # 尝试导入 Detectron2LayoutModel
        try:
            from layoutparser.models import Detectron2LayoutModel
            print("LayoutParser 完整安装")
        except ImportError:
            print("正在安装 LayoutParser 的额外依赖...")
            subprocess.check_call([
                sys.executable, "-m", "pip", "install", "layoutparser[layoutmodels]"
            ])
    except ImportError:
        print("正在安装 LayoutParser...")
        subprocess.check_call([
            sys.executable, "-m", "pip", "install", "layoutparser"
        ])
        print("正在安装 LayoutParser 的额外依赖...")
        subprocess.check_call([
            sys.executable, "-m", "pip", "install", "layoutparser[layoutmodels]"
        ])
    
    # 检查 torch 和 torchvision
    try:
        import torch
        import torchvision
        print(f"PyTorch 版本: {torch.__version__}")
        print(f"Torchvision 版本: {torchvision.__version__}")
    except ImportError:
        print("正在安装 PyTorch 和 Torchvision...")
        subprocess.check_call([
            sys.executable, "-m", "pip", "install", "torch", "torchvision"
        ])

def update_imageocr_script():
    """更新 imageocr.py 脚本以修复导入问题"""
    try:
        with open('imageocr.py', 'r', encoding='utf-8') as f:
            content = f.read()
        
        # 修改导入语句
        updated_content = content.replace(
            "import layoutparser as lp", 
            "import layoutparser as lp\nfrom layoutparser.models import Detectron2LayoutModel"
        )
        
        # 修改模型初始化
        updated_content = updated_content.replace(
            "self.layout_model = lp.Detectron2LayoutModel(",
            "self.layout_model = Detectron2LayoutModel("
        )
        
        with open('imageocr.py', 'w', encoding='utf-8') as f:
            f.write(updated_content)
        
        print("已更新 imageocr.py 脚本")
    except Exception as e:
        print(f"更新脚本时出错: {e}")

if __name__ == "__main__":
    print("开始修复 LayoutParser 问题...")
    check_and_install_dependencies()
    update_imageocr_script()
    print("\n修复完成！请重新运行 imageocr.py") 